#include <iostream>
#include <vector>
#include <thread>
#include <queue>
#include <functional>
#include <mutex>
#include <condition_variable>
#include <future>
#include <atomic>

class ThreadPool {
public:
    // 构造函数：创建指定数量的线程
    explicit ThreadPool(size_t num_threads) : stop(false) {
        for (size_t i = 0; i < num_threads; ++i) {
            workers.emplace_back([this] {
                for (;;) { // 工作循环
                    std::function<void()> task;

                    {
                        // 获取任务锁
                        std::unique_lock<std::mutex> lock(queue_mutex);
                        // 等待条件变量，直到有任务或线程池停止
                        condition.wait(lock, [this] { return stop || !tasks.empty(); });

                        // 如果线程池已停止且任务队列为空，则退出循环
                        if (stop && tasks.empty()) {
                            return;
                        }

                        // 从队列中取出一个任务
                        task = std::move(tasks.front());
                        tasks.pop();
                    }

                    // 执行取出的任务
                    task();
                }
            });
        }
    }

    // 析构函数：等待所有任务完成并停止所有线程
    ~ThreadPool() {
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            stop = true;
        }
        condition.notify_all(); // 通知所有等待的线程

        for (std::thread &worker : workers) {
            worker.join(); // 等待每个线程完成
        }
    }

    // 模板方法：提交任务到线程池
    // 支持任何可调用对象（函数、lambda、函数对象等）以及参数
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args) -> std::future<typename std::result_of<F(Args...)>::type> {
        // 推导任务函数的返回类型
        using return_type = typename std::result_of<F(Args...)>::type;

        // 创建一个打包了任务和参数的 std::function
        auto task = std::make_shared<std::packaged_task<return_type()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...)
        );

        // 获取返回值的 future
        std::future<return_type> res = task->get_future();

        {
            // 获取任务队列锁
            std::unique_lock<std::mutex> lock(queue_mutex);

            // 如果线程池已经停止，则抛出异常
            if (stop) {
                throw std::runtime_error("enqueue on stopped ThreadPool");
            }

            // 将任务添加到队列中
            tasks.emplace([task] { (*task)(); });
        }

        condition.notify_one(); // 通知一个等待的线程来执行任务
        return res;
    }

private:
    // 工作线程容器
    std::vector<std::thread> workers;
    // 任务队列
    std::queue<std::function<void()>> tasks;

    // 用于同步的互斥锁和条件变量
    std::mutex queue_mutex;
    std::condition_variable condition;
    // 原子布尔值，指示线程池是否应该停止
    std::atomic<bool> stop;
};

// 示例用法
int main() {
    const int num_threads = 4; // 创建一个包含4个线程的线程池
    ThreadPool pool(num_threads);

    std::vector<std::future<int>> results; // 存储任务返回的 future

    // 提交10个任务到线程池
    for (int i = 0; i < 10; ++i) {
        results.emplace_back(
            pool.enqueue([i] {
                std::cout << "Task " << i << " is running on thread " << std::this_thread::get_id() << std::endl;
                // 模拟一些工作
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                std::cout << "Task " << i << " finished." << std::endl;
                return i * i; // 返回 i 的平方
            })
        );
    }

    // 获取每个任务的返回值
    for (auto&& result : results) {
        std::cout << "Result: " << result.get() << std::endl;
    }

    std::cout << "All tasks completed." << std::endl;

    // main 函数结束时，ThreadPool 的析构函数会被自动调用，
    // 它会等待所有剩余任务完成并停止所有工作线程。
    return 0;
}